import streamlit as st
from webui_pages.utils import *
from streamlit_chatbox import *
from datetime import datetime
from typing import List, Dict
import os
import copy
import itertools

chat_box = ChatBox(
    assistant_avatar=os.path.join(
        "img",
        "chatchat_icon_blue_square_v2.png"
    ),
    session_key="file_agent"
)


def get_messages_history(history_len: int) -> List[Dict]:
    def filter(msg):
        '''
        针对当前简单文本对话，只返回每条消息的第一个element的内容
        '''
        content = [x._content for x in msg["elements"] if x._output_method in ["markdown", "text"]]
        return {
            "role": msg["role"],
            "content": content[0] if content else "",
        }

    history = chat_box.filter_history(100000, filter)  # workaround before upgrading streamlit-chatbox.
    user_count = 0
    i = 1
    for i in range(1, len(history) + 1):
        if history[-i]["role"] == "user":
            user_count += 1
            if user_count >= history_len:
                break
    return history[-i:]


# shy-2023-09-06
# 功能都是通过调用api对象的方法，来实现的。
def file_agent_page(api: ApiRequest, server_config):
    chat_box.init_session()

    with st.sidebar:
        def on_mode_change():
            mode = st.session_state.dialogue_mode
            text = f"已切换到 {mode} 模式。"
            st.toast(text)
            # sac.alert(text, description="descp", type="success", closable=True, banner=True)

        dialogue_mode = st.selectbox("请选择对话模式",
                                     ["LLM 对话",
                                      ],
                                     on_change=on_mode_change,
                                     key="dialogue_mode",
                                     )

    docs_present = api.list_docs_in_file_agent()
    if not len(docs_present):
        st.info(f"当前文档:无")
    else:
        st.info(f"当前文档:{docs_present}")

    if st.button(
            "清空当前文档",
            disabled=len(docs_present) == 0,
    ):
        api.delete_docs_in_file_agent()
        st.experimental_rerun()

    files = st.file_uploader("上传目标文档",
                             ['csv'],
                             accept_multiple_files=True,
                             )

    if st.button(
            "确定上传",
            disabled=len(files) == 0,
    ):
        data = [{"file": f} for f in files]
        # data[-1]["not_refresh_vs_cache"]=False
        for kwargs in data:
            ret = api.upload_files_for_agent(**kwargs)
            if msg := check_success_msg(ret):
                st.toast(msg, icon="✔")
            elif msg := check_error_msg(ret):
                st.toast(msg, icon="✖")

        ret = api.create_agent_for_files()
        if msg := check_success_msg(ret):
            st.toast(msg, icon="✔")
        elif msg := check_error_msg(ret):
            st.toast(msg, icon="✖")

        st.session_state.files = []
        st.experimental_rerun()

    st.divider()

    # Display chat messages from history on app rerun

    chat_box.output_messages()

    chat_input_placeholder = "请输入对话内容，换行请使用Ctrl+Enter "

    # !!!streamlit文档说Chat input can only be used once per app page，否则会多线程冲突，
    # 从而页面会不断刷新造成st.session_state为{}!!!!!!!!!!!!

    # 这个key很重要，若同一个app page出现2个key，则相当于chat input被用了2次。
    prompt = st.chat_input(chat_input_placeholder, key="prompt")
    if prompt:
        # 若是使用stream=True，则会返回一个迭代器
        # 若是使用stream=False，则会返回一个字符串
        use_stream = True
        if use_stream:
            r = api.ask_file_agent(prompt, stream=use_stream)

            if r:

                text = ''
                # 通过try来捕捉那些st.session_state为{}的线程的报错，pass掉。
                try:

                    for t in r:
                        chat_box.user_say(prompt)
                        chat_box.ai_say("正在思考...")

                        if error_msg := check_error_msg(t):  # check whether error occured
                            st.error(error_msg)
                            break
                        text += t
                        chat_box.update_msg(text)

                    chat_box.update_msg(text, streaming=False)  # 更新最终的字符串，去除光标
                except Exception:
                    pass

        else:

            text = api.ask_file_agent(prompt, stream=use_stream)
            # 通过try来捕捉那些st.session_state为{}的线程的报错，pass掉。
            try:
                chat_box.user_say(prompt)
                chat_box.ai_say("正在思考...")
                if text:
                    chat_box.update_msg(text, streaming=False)  # 更新最终的字符串，去除光标
            except Exception as e:
                pass

    now = datetime.now()
    with st.sidebar:

        cols = st.columns(2)
        export_btn = cols[0]
        if cols[1].button(
                "清空对话",
                use_container_width=True,
        ):
            chat_box.reset_history()
            st.experimental_rerun()

    export_btn.download_button(
        "导出记录",
        "".join(chat_box.export2md()),
        file_name=f"{now:%Y-%m-%d %H.%M}_对话记录.md",
        mime="text/markdown",
        use_container_width=True,
    )
